#!/usr/bin/python3
import os, sys
import datetime, fileinput, multiprocessing
from bed import *
from File import FileWriter


def about():
    print('\t\t\t\t\t本脚本')
    print('\t用来搜索两个bed文件的具有重合的部分')
    print('\t输出文件为bed6格式的文件\n')
    print('\t##########\t文档说明\t##########')
    print('\t-i：要进行比对的文件，单个文件每次')
    print('\t-o：要输出的文件，单个文件每次')
    print('\t-g：参考文件，所有bed最好经过排序')
    print('\t-p：要使用的进程数')
    # print('\t注意：算法永远比语种更关键')
    exit()


class RealHit(multiprocessing.Process):
    def __init__(self, chrlist, mybed, refbed, thread, outputfile, multi=True):
        multiprocessing.Process.__init__(self)
        self.chrlist = chrlist
        self.mybed = mybed
        self.refbed = refbed
        self.thread = thread
        self.outputfile = outputfile
        self.multi = multi

    def run(self):
        chrlist = self.chrlist
        mybed = self.mybed
        refbed = self.refbed
        thread = self.thread
        outputfile = self.outputfile

        if self.multi:
            outputfile = outputfile + '.tem' + str(thread)

        if os.path.isfile(outputfile):
            os.remove(outputfile)

        position_my = get_position(mybed)
        position_ref = get_position(refbed)
        chrlist.sort()
        have = []
        writed = set()
        with FileWriter(outputfile) as w:
            for i in range(len(chrlist)):                       # 根据不同染色体在bed文件中的分布不同，分别提取出来，单独跑，大量时间
                if i < len(chrlist)-1:
                    mybed_tem = mybed[position_my.get(chrlist[i]):(position_my.get(chrlist[i+1])-1)]
                    refbed_tem = refbed[position_ref.get(chrlist[i]):(position_ref.get(chrlist[i+1])-1)]
                else:
                    mybed_tem = mybed[position_my.get(chrlist[i]):]
                    refbed_tem = refbed[position_ref.get(chrlist[i]):]

                for my in mybed_tem:
                    for ref in refbed_tem:
                        if my & ref:                             # 重载了bed类中的__and__方法，正好可以用来判断两者是否有重叠
                            if my.getAll() not in writed:
                                writed.add(my.getAll())          # 一个set记录是否已经输出过这个文件了，如果需搜索两个结合位点数据
                                w.write(my.getAll() + "\n")      # 就在里进行修改
                                if my.getChromosome() not in have:          # 这个很简单，就是为了输出进度而已
                                    have.append(my.getChromosome())
                                    print("线程" + str(thread) + "：统计" + my.getChromosome() + "中")
                            else:
                                break


def divide_list(chrlist, thread):                                   # 将list文件根据进程数进行分割
    size = len(chrlist) // thread
    results = []
    for i in range(thread):
        if i < (thread - 1):
            results.append(chrlist[(size * i): size * (i + 1)])
        else:
            results.append(chrlist[size * i:])
    return results


def get_position(bed):                                              # 统计出每个bed文件不同染色体在list中的起点
    position = {}

    for i in range(len(bed)):
        if bed[i].getChromosome() in position.keys():
            continue
        else:
            position.update({bed[i].getChromosome(): i})
    return position


def get_hit(inputfile, reference, outputfile, thread):
    begin = datetime.datetime.now()
    print("程序开始于：%s" % begin)
    try:
        # 首先读取文件
        mybed = []
        myList = set()
        with fileinput.input(inputfile) as f:
            for line in f:
                lines = re.split("\\s+", line)
                if len(lines) < 6:
                    print("input pass")
                    continue
                else:
                    myList.add(lines[0])
                    # 将要比对的bed读入
                    mybed.append(bed(chromosome=lines[0], start=int(lines[1]), end=int(lines[2]),
                                     name=lines[3], score=int(lines[4]), strand=lines[5]))
        mybed.sort()
        print("%s读取完成" % inputfile)

        # 读取ref文件，其实这样子并不节省内存，但是为了省下sort的麻烦可以了
        # 后期有心情可以多写一个模块，单独用来sort bed这种文件
        refbed = []
        refList = set()
        with fileinput.input(reference) as f:
            for line in f:
                lines = re.split("\\s+", line)
                if len(lines) < 6:
                    print("reference pass")
                    continue
                else:
                    if lines[0] in myList:                                  # 如果第一个文件没有的chr数据就不读取了，没有意义
                        refList.add(lines[0])
                        # 将要比对的bed读入
                        refbed.append(bed(chromosome=lines[0], start=int(lines[1]), end=int(lines[2]),
                                         name=lines[3], score=int(lines[4]), strand=lines[5]))
        refbed.sort()
        print("%s读取完成" % reference)
        if len(refbed) == 0:
            print('reference is empty')
        elif len(mybed) == 0:
            print('input file is empty')
        chrlist = list(myList.intersection(refList))                          # 提取两个bed文件相同的chr，节省后边的计算
        # 清除无用变量节省内存
        myList = None
        refList = None

        if len(chrlist) < thread:
            thread = len(chrlist)

    except BedFormatError as e:
        print(e.value)

    if thread == 1:
        proc = RealHit(chrlist=chrlist, mybed=mybed, refbed=refbed, thread=thread, outputfile=outputfile, multi=False)
        proc.start()
        proc.join()
    else:
        divided_list = divide_list(chrlist, thread)
        proc = []
        for i in range(thread):
            proc.append(RealHit(chrlist=divided_list[i], mybed=mybed, refbed=refbed,
                                thread=(i + 1), outputfile=outputfile, multi=True))

        for i in range(thread):
            proc[i].start()

        for i in range(thread):
            proc[i].join()

        have = set()
        with FileWriter(outputfile, mode='w') as w:
            for i in range(thread):
                temFile = outputfile + '.tem' + str(i + 1)
                if os.path.isfile(temFile):
                    with fileinput.input(temFile) as f:
                        for line in f:
                            if line not in have:
                                w.write(line)                                   # 将数据从临时文件，转移到最终文件中
                                have.add(line)                                  # 在这个过程中还判断是否有重复的数据
                else:
                    print("临时文件%s不存在" % temFile)
                    exit()
                os.remove(temFile)                                          # 临时文件用完以后，删掉
    end = datetime.datetime.now()
    print("程序结束于：%s" % end)
    print("共计用时：%s" % str(end - begin))

if __name__ == '__main__':
    parameters = {}  # 用来读取命令行参数的，遍历所有参数，开头以-开始的，下一个就作为参数附加，然后统一添加到一个词典中，再提取出来
    if len(sys.argv) <= 1:
        about()
    try:
        for i in range(len(sys.argv)):
            if sys.argv[i] == '-h' or sys.argv[i] == '--help':
                about()
            if sys.argv[i].startswith('-'):
                parameters.update({sys.argv[i]: sys.argv[i + 1]})
            elif sys.argv[i].startswith('--'):
                parameters.update({sys.argv[i]: sys.argv[i + 1]})
    except IndexError:  # 判断参数是否输入有误
        print("输入的参数有误")
        exit()

    inputfile = parameters.get('-i')                         # 提取所有的输入文件

    if inputfile:                                            # 判断一下输入的文件是否是真的文件
        if not os.path.isfile(inputfile):
            print(inputfile + "不是文件")
            exit()
    else:
        print("没有任何输入文件")
        exit()

    outputfile = parameters.get('-o')                       # 提取所有的输出文件
    outputDir = None

    if outputfile:                                          # 判断是否输入了输出文件
        if os.path.isfile(outputfile):                      # 如果文件存在，删除，其实这个不要紧，可以直接w覆盖，这个是为了保险
            os.remove(outputfile)
    else:
        print('没有指定任何输出文件')

    if not parameters.get('-p'):                            # 判断进程数，默认为1
        thread = 1
    else:
        try:

            thread = int(parameters.get('-p'))
            if thread < 1:                                  # thread 如果没有设置之类的，默认为1
                thread = 1
        except ValueError:
            print("-p参数需要输入整数")

    refFile = parameters.get('-g')

    if refFile:                                            # 判断一下输入的文件是否是真的文件
        if not os.path.isfile(refFile):
            print(refFile + "不是文件")
            exit()
    else:
        print("没有指定任何参考文件")
        exit()

    # ！！！！！！！！！！！！！！！！！！！！！！！！！！！！！       参数处理       ！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！
    get_hit(inputfile=inputfile, outputfile=outputfile, thread=thread, reference=refFile)
    input("任意键退出")
